#ifndef RB_TREE_H_
#define RB_TREE_H_

#include <vector>
#include <time.h> 
#include <memory> 

using namespace std;
/*
 * @brief red-black tree
 * 参考: https://blog.csdn.net/cy973071263/article/details/122543826
 */
typedef short HEIGHT_TYPE;

template<class T>
class RBTreeNode{
protected:
	T val;
	shared_ptr<RBTreeNode<T>> children[2]; 
	weak_ptr<RBTreeNode<T>> father; //父节点
	bool isRed = true; //是否红色节点
	HEIGHT_TYPE height;
	template<typename> friend class RBTree;

public:
	RBTreeNode(T val){
		this->val = val;	
		children[0] = children[1] = nullptr;
		height = 1;
	}
	
	HEIGHT_TYPE getHeight(){
		return height;
	}
	
	HEIGHT_TYPE balanceFactor(){
		HEIGHT_TYPE factor = 0;
		if(children[0] != nullptr){
			factor += children[0]->height;
		}
		if(children[1] != nullptr){
			factor -= children[1]->height;
		}
		return factor;
	}

	//更新树的高度
	void updateHeight(){
		HEIGHT_TYPE h = 0;	
		if(children[0] != nullptr && children[0]->height > h){
				h = children[0]->height;
		}
		if(children[1] != nullptr && children[1]->height > h){
				h = children[1]->height;
		}
		height = h + 1;
	}

	void traversal(vector<RBTreeNode<T>> &nodes){
		if(children[0] != nullptr){
			children[0]->traversal(nodes);
		}
		nodes.push_back(*this);
		if(children[1] != nullptr){
			children[1]->traversal(nodes);
		}
	}

	T getVal(){
		return val;
	}

	bool isRedNode(){
        return isRed;
    }

	bool isBlackNode(){
        return !isRed;
    }

	void setRed(){
        isRed = true;
    }

	void setBlack(){
        isRed = false;
    }

    //获取节点是父节点的第几个子节点
	int childIndex(){
        int index = -1;
        auto fatherNode = father.lock();
        if(fatherNode != nullptr){
            for(int i = 0; i < 2; i++){
                if(fatherNode->children[i].get() == this){
                    index = i;
                    break;
                }
            }
        }
        return index;
    }

    //获取叔叔节点
	shared_ptr<RBTreeNode<T>> getUncle(){
        auto fatherNode = father.lock();
        if(fatherNode == nullptr){
            return nullptr;
            
        }
        auto grandfather = fatherNode->father.lock();
        if(grandfather == nullptr){
            return nullptr;
            
        }
        for(int i = 0; i < 2; i++){
            if(grandfather->children[i].get() != fatherNode.get()){
                return grandfather->children[i];
            }
        }
        return nullptr;
    }

    //孩子节点个数
	int childNum(){
        int num = 0;
        for(int i = 0; i < 2; i++){
            if(children[i] != nullptr){
                num++;
            }
        }
        return num;
    }

};

template<class T>
class RBTree{
protected:
	shared_ptr<RBTreeNode<T>> root;	

	//重新平衡树
	void insertFix(shared_ptr<RBTreeNode<T>> node){
		shared_ptr<RBTreeNode<T>> x = node;
        int in, ip;
		while(x != nullptr){
            auto fatherNode = x->father.lock();
            if(fatherNode == nullptr){
                // x is root node
                x->setBlack();
                break;
            }
            auto grandfather = fatherNode->father.lock();
            if( grandfather == nullptr || fatherNode->isBlackNode()){
                break;
            }
            //父节点是红色
            in = x->childIndex();
            ip = fatherNode->childIndex();
            auto uncle = x->getUncle();
            if(uncle == nullptr || uncle->isBlackNode()){
                //has no uncle or uncle is black node
                if(in == ip){
                    //LL or RR insert
                    fatherNode->setBlack();
                    grandfather->setRed();
                    rotate(grandfather, ip ^ 1);
                }else{
                    //LR or RL insert
                    node->setBlack();
                    grandfather->setRed();
                    rotate(fatherNode, in ^ 1);
                    rotate(grandfather, ip ^ 1);
                }
                break;
            }else{
                //has uncle and uncle is red node
                //LL, RR, LR or RL insert
                grandfather->setRed();
                fatherNode->setBlack();
                uncle->setBlack();
                x = grandfather;
            }
		}
	}

public:

	RBTree(){
		root = nullptr;
	}

	void insert(T val){
		shared_ptr<RBTreeNode<T>> node = make_shared<RBTreeNode<T>>(val);
		shared_ptr<RBTreeNode<T>> cur = root, p = nullptr;
		while(cur != nullptr){
			p = cur;
			if(val < cur->getVal()){
				cur = cur->children[0];
			}else{
				cur = cur->children[1];
			}
		}
		if(p == nullptr){
			root = node;
            root->isRed = false; //根节点为黑色节点
		}else{
			node->father = p;
			node->isRed = true;
			if(val < p->getVal()){
				p->children[0] = node;
			}else{
				p->children[1] = node;
			}
            if(node->father.lock()->isRedNode()){
			    insertFix(node);
            }
		}
	}
    
	// index = 0 for left rotate , index = 1 for right rotate
	void rotate(shared_ptr<RBTreeNode<T>> x, int index){
        int otherIndex = index ^ 1;
		shared_ptr<RBTreeNode<T>> y = x->children[otherIndex];
		if(y == nullptr) return ;
		if(x->father.lock() != nullptr){
			if(x->father.lock()->children[index] == x){
				x->father.lock()->children[index] = y;
			}else{
				x->father.lock()->children[otherIndex] = y;
			}
		}else{
			root = y;
		}
		y->father = x->father;
		x->father = y;
		x->children[otherIndex] = y->children[index];
		if(y->children[index] != nullptr){
			y->children[index]->father = x;
		}
		y->children[index] = x;
		//调整树的高度
		x->updateHeight();
		y->updateHeight();
		if(y->father.lock() != nullptr){
			y->father.lock()->updateHeight();
		}
	}

	//遍历树
	vector<RBTreeNode<T>> traversal(){
		vector<RBTreeNode<T>> nodes;
		if(root != nullptr){
			root->traversal(nodes);
		}
		return nodes;
	}

	//随机生成一颗树
	void intRandomTree(int n){
		vector<int> nums(n);
		for(int i = 0; i < n; i++){
			nums[i] = i + 1;
		}
		srand((unsigned)time(nullptr));
		int pos;
		for(int i = n - 1; i >= 0; i--){
			pos = random() % (i + 1);
			insert(nums[pos]);
			nums[pos] = nums[i];
		}
	}

	//打印树
	void printTree(){
		cout << "{" << endl;
		if(root != nullptr){
			vector<shared_ptr<RBTreeNode<T>>> que;			
			vector<int> 	depth;
			shared_ptr<RBTreeNode<T>> head;
			que.push_back(root);
			depth.push_back(1);
			int l, d;
			l = 0;
			// BFS遍历
			while(l < que.size()){
				head = que[l];
				d = depth[l];
				if(head->children[0] != nullptr){
					que.push_back(head->children[0]);
					depth.push_back(d + 1);
				}
				if(head->children[1] != nullptr){
					que.push_back(head->children[1]);
					depth.push_back(d + 1);
				}
				l++;	
			}
			d = 1;
			for(int i = 0; i < que.size(); i++){
				if(i > 0){
					if(depth[i] != d){
						cout << endl;
					}else{
						cout << " ";
					}
					d = depth[i];
				}
				cout << "(" ;
				if(que[i]->father.lock() != nullptr){
					cout << que[i]->father.lock()->getVal();	
				}else{
					cout << "nullptr";	
				}
				cout << "," << que[i]->getVal();
                if(que[i]->isRedNode()){
                    cout << "," << "r";
                }else{
                    cout << "," << "b";
                }
				cout << ")";
			}
		}
		cout << endl << "}" << endl;
	}

	//按递增顺序打印
	void printSerial(){
		vector<RBTreeNode<int>> nodes = traversal();
		cout << "[" ;
		for(int i = 0; i < nodes.size(); i++){
			if(i > 0){
				cout <<",";
			}
			cout << nodes[i].getVal();
		}
		cout << "]";
	}

	// tangwulog 测试 insert函数
	static void test(){
        RBTree<int> rbTree;
        rbTree.insert(3);
        rbTree.insert(1);
        rbTree.insert(2);
        rbTree.printTree();
    }

	//批量测试
	static void batchTest(){
		int batchNum = 1;
        int num = 2;
		for(int i = 0; i < batchNum; i++){
			RBTree<int> rbTree;
			rbTree.intRandomTree(num);
			vector<RBTreeNode<int>> nodes = rbTree.traversal();
			bool flag = true;
            cout << "rbtree {" << nodes[0].getVal() << "";
			for(int j = 1; j < nodes.size(); j++){
				if(nodes[j-1].getVal() > nodes[j].getVal()){
					flag = false;
					break;
				}
                cout << ", " << nodes[j].getVal();
			}
            cout << "}" << endl;
			if(flag){
				cout << "index=" << i << ",success" << endl;
			}else{
				cout << "index=" << i << ",error ";
				rbTree.printSerial();
				cout << endl;
			}
		}
	}

	shared_ptr<RBTreeNode<T>> getRoot(){
		return root;
	}

};

#endif
